import cv2
import os
import numpy as np

data_path = './train_images'

# 使用OpenCV用来检测脸部的函数
def detect_face(img):
    # 将测试图像转换为灰度图像，因为opencv人脸检测器需要灰度图像
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5)

    # 如果未检测到面部，则返回原始图像
    if (len(faces) == 0):
        return None,None

    (x, y, w, h) = faces[0]
    # 只返回图像的正面部分
    return gray[y:y + w, x:x + h], faces[0]

def prepare_training_data(data_folder_path):
    dirs = os.listdir(data_folder_path)
    # os.listdir可以获取当前路径下一级所有的子文件路径（返回值是一个字符串数组）
    faces = []
    labels = []
    for dir_name in dirs:
        label = int(dir_name)
        # sample subject_dir_path = "train_images/0"
        subject_dir_path = data_folder_path + '/' + dir_name
        subject_images_names = os.listdir(subject_dir_path)
        # print(subject_images_names)
        for image_name in subject_images_names:
            # 忽略.DS_Store之类的系统文件
            if image_name.startswith("."):
                continue
            # sample image path = train_images/0/1.jpg
            # print('dir_name:',dir_name)
            image_path = subject_dir_path + '/' + image_name
            print('The path of training image:',image_path)
            # 阅读图像
            # image_path_ = image_path.replace('0','',1)
            image = cv2.imread(image_path)
            face, rect = detect_face(image)
            if face is not None:
                # 将脸添加到脸部列表
                faces.append(face)
                # 为这张脸添加标签
                labels.append(label)
    return faces, labels


if __name__ == '__main__':
    print('Training faces. It will take a few seconds. Wait ...')
    faces, labels = prepare_training_data(data_path)
    face_recognizer = cv2.face.LBPHFaceRecognizer_create()
    face_recognizer.train(faces, np.array(labels))
    face_recognizer.write(r'./train_data.yml')
    print('The data of faces has been saved...')